import torch
import numpy as np
import matplotlib.pyplot as plt
from torch import nn  # 补充缺失的nn模块导入

from lib.modeling.basic_blocks import local_p_block, single_p_block, global_p_block


class HGLM_Visualizer:
    def __init__(self, part_num=6, in_channels=256, out_channels=256):
        self.part_num = part_num
        self.in_channels = in_channels
        self.out_channels = out_channels

        # 初始化支路
        self.s_path = single_p_block(part_num, in_channels, out_channels)
        self.l_path = local_p_block(part_num, in_channels, out_channels)
        self.g_path = global_p_block(part_num, in_channels, out_channels)

        # 修复权重生成器
        self.weight_generator = nn.Sequential(
            nn.Linear(3, 3),  # 输入维度调整为3
            nn.Softmax(dim=1)
        )

    def _generate_dummy_input(self, batch_size=2, seq_len=8):
        """生成输入特征 (p, n, c, s)"""
        return torch.randn(
            self.part_num,
            batch_size,
            self.in_channels,
            seq_len
        )

    def _compute_attention_weights(self, features):
        """修复后的权重计算"""
        # 计算每个支路的均值特征 [batch_size, 1]
        pooled = [f.mean(dim=(0, 2, 3)).unsqueeze(1) for f in features]
        # 沿dim=1拼接 -> [batch_size, 3]
        pooled_features = torch.cat(pooled, dim=1)
        return self.weight_generator(pooled_features)

    def visualize(self, save_path="hglm_visualization.png"):
        # 生成输入数据
        x = self._generate_dummy_input()

        # 前向传播各支路
        s_out = self.s_path(x)
        l_out = self.l_path(x)
        g_out = self.g_path(x)

        # 计算动态权重
        weights = self._compute_attention_weights([s_out, l_out, g_out])

        # 特征降维（示例取通道均值）
        def reduce_dim(feature):
            return feature.mean(dim=(2, 3)).cpu().detach().numpy()

        # 准备可视化数据
        vis_data = {
            'single': reduce_dim(s_out),
            'local': reduce_dim(l_out),
            'global': reduce_dim(g_out),
            'weights': weights.detach().numpy()
        }

        # 创建画布
        fig, axs = plt.subplots(2, 2, figsize=(12, 10))

        # 绘制单点支路特征热力图
        im1 = axs[0, 0].imshow(vis_data['single'], cmap='viridis', aspect='auto')
        axs[0, 0].set_title("Single-Path Features")
        fig.colorbar(im1, ax=axs[0, 0])

        # 绘制局部支路特征热力图
        im2 = axs[0, 1].imshow(vis_data['local'], cmap='plasma', aspect='auto')
        axs[0, 1].set_title("Local-Path Features")
        fig.colorbar(im2, ax=axs[0, 1])

        # 绘制全局支路特征热力图
        im3 = axs[1, 0].imshow(vis_data['global'], cmap='inferno', aspect='auto')
        axs[1, 0].set_title("Global-Path Features")
        fig.colorbar(im3, ax=axs[1, 0])

        # 绘制动态权重分布
        labels = ['Single', 'Local', 'Global']
        axs[1, 1].bar(labels, vis_data['weights'][0], color=['#1f77b4', '#ff7f0e', '#2ca02c'])
        axs[1, 1].set_ylim(0, 1)
        axs[1, 1].set_title("Dynamic Fusion Weights")

        # 保存图像
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()


if __name__ == "__main__":
    visualizer = HGLM_Visualizer(part_num=6, in_channels=256, out_channels=256)
    visualizer.visualize(save_path="hglm_fusion_visualization.png")